import matplotlib.pyplot as plt
import numpy as np

# plot params
fontsize = 16
markersize = 8
linestyles = ['-', '--', ':', '-.', '-']
markers = ['^', '+', 'p', 's', '<', '>', '*', 'x']
colors = [
    '#0C5DA5', '#00B945', '#FF9500', '#FF2C00', '#845B97', '#474747', '#9e9e9e'
]  # scienceplot color palette

plt.style.use(['science', 'no-latex'])
plt.rcParams.update({
    "font.family": "sans-serif",
    "font.serif": ["arial"],
    "font.size": fontsize,
    'ps.fonttype': 42
})

# Model names and data
models = ["Q-learning", "FNN", "Expert (MLP)", "Mr.CL"]
flops = [0, 1163.0, 8768.0, 27872.0]  # FLOPs
params = [2620104.8, 1220.0, 8964.0, 28527.0]  # Parameters (Q-learning has NA)

# Creating the figure and axes
fig, ax = plt.subplots(figsize=(7, 5), dpi=200)

# Bar positions
ind = np.arange(len(models))
width = 0.35  # Bar width

# Plot FLOPs and Parameters as separate bars
bars_flops = ax.bar(ind - width / 2, flops, width, label='FLOPs', facecolor=colors[2], hatch="-")
bars_params = ax.bar(ind + width / 2, params, width, label='Parameters', facecolor=colors[3], hatch="x")

# Set a logarithmic scale on the y-axis to handle large Q-learning values better
ax.set_yscale('log')

# Adding labels and grid
ax.tick_params(axis='x', direction='in', bottom=True)
ax.tick_params(axis='y', direction='out', labelsize=fontsize, length=3)
ax.set_xticks(ind)
ax.set_xticklabels(models, size=fontsize)
ax.set_xlabel('Models', size=fontsize)
ax.set_ylabel('FLOPs / Parameters (log scale)', size=fontsize)
ax.xaxis.grid(which='major', linestyle='--')
ax.yaxis.grid(which='major', linestyle='--')

# Adding the legend
ax.legend(prop={'size': fontsize}, ncol=2, loc='upper center', bbox_to_anchor=(0.5, 1.15))

# Save and show the figure
fig.savefig('./model_comparison_flops_params_log.eps', bbox_inches='tight')
plt.show()
